// Epart.cpp : Implementation of DLL Exports.

//
// Note: COM+ 1.0 Information:
//      Please remember to run Microsoft Transaction Explorer to install the component(s).
//      Registration is not done by default. 

#include "stdafx.h"
#include "resource.h"
#include "Epart_i.h"
#include "dllmain.h"
#include "compreg.h"
#include "aes.h"
#include "md5.h"
#include <iostream>
#include <vector>
#include <fstream>
#include <exception>
#include <string>
#include <comdef.h>
#include <stdio.h>
using namespace std;
typedef unsigned int uint;

// Used to determine whether the DLL can be unloaded by OLE.
STDAPI DllCanUnloadNow(void)
{
			return _AtlModule.DllCanUnloadNow();
	}

// Returns a class factory to create an object of the requested type.
STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID* ppv)
{
		return _AtlModule.DllGetClassObject(rclsid, riid, ppv);
}

// DllRegisterServer - Adds entries to the system registry.
STDAPI DllRegisterServer(void)
{
	// registers object, typelib and all interfaces in typelib
	HRESULT hr = _AtlModule.DllRegisterServer();
		return hr;
}

// DllUnregisterServer - Removes entries from the system registry.
STDAPI DllUnregisterServer(void)
{
	HRESULT hr = _AtlModule.DllUnregisterServer();
		return hr;
}

// DllInstall - Adds/Removes entries to the system registry per user per machine.
STDAPI DllInstall(BOOL bInstall, LPCWSTR pszCmdLine)
{
	HRESULT hr = E_FAIL;
	static const wchar_t szUserSwitch[] = L"user";

	if (pszCmdLine != NULL)
	{
		if (_wcsnicmp(pszCmdLine, szUserSwitch, _countof(szUserSwitch)) == 0)
		{
			ATL::AtlSetPerUserRegistration(true);
		}
	}

	if (bInstall)
	{	
		hr = DllRegisterServer();
		if (FAILED(hr))
		{
			DllUnregisterServer();
		}
	}
	else
	{
		hr = DllUnregisterServer();
	}

	return hr;
}


// Epart.cpp : Implementation of CEpart

#include "stdafx.h"
#include "Epart.h"


// CEpart


STDMETHODIMP CEpart::OnStartPage (IUnknown* pUnk)  
{
	if(!pUnk)
		return E_POINTER;

	CComPtr<IScriptingContext> spContext;
	HRESULT hr;

	// Get the IScriptingContext Interface
	hr = pUnk->QueryInterface(__uuidof(IScriptingContext), (void **)&spContext);
	if(FAILED(hr))
		return hr;

	// Get Request Object Pointer
	hr = spContext->get_Request(&m_piRequest);
	if(FAILED(hr))
	{
		return hr;
	}

	// Get Response Object Pointer
	hr = spContext->get_Response(&m_piResponse);
	if(FAILED(hr))
	{
		m_piRequest.Release();
		return hr;
	}
	
	// Get Server Object Pointer
	hr = spContext->get_Server(&m_piServer);
	if(FAILED(hr))
	{
		m_piRequest.Release();
		m_piResponse.Release();
		return hr;
	}
	
	// Get Session Object Pointer
	hr = spContext->get_Session(&m_piSession);
	if(FAILED(hr))
	{
		m_piRequest.Release();
		m_piResponse.Release();
		m_piServer.Release();
		return hr;
	}

	// Get Application Object Pointer
	hr = spContext->get_Application(&m_piApplication);
	if(FAILED(hr))
	{
		m_piRequest.Release();
		m_piResponse.Release();
		m_piServer.Release();
		m_piSession.Release();
		return hr;
	}
	m_bOnStartPageCalled = TRUE;
	return S_OK;
}

STDMETHODIMP CEpart::OnEndPage ()  
{
	m_bOnStartPageCalled = FALSE;
	// Release all interfaces
	m_piRequest.Release();
	m_piResponse.Release();
	m_piServer.Release();
	m_piSession.Release();
	m_piApplication.Release();
	return S_OK;
}

/*!
 * This method sets a secret key string.
 *
 * \param[in] bstrNewVal	A key string which is converted into UTF-8 and hashed by MD5.
 *                          Null or an empty string is not allowed.
 */
STDMETHODIMP CEpart::put_Key(const BSTR bstrNewVal)
{
	if ( bstrNewVal == NULL )
	{
		_Error (OLESTR("Key: The key can not be null.")) ;
		return E_INVALIDARG ;
	}
	
	if ( ::SysStringLen (bstrNewVal) == 0 )
	{
		_Error (OLESTR("Key: The key can not be an empty string.")) ;
		return E_INVALIDARG ;
	}

	// If the new key and the old key are the same
	if ( (m_bstrKey != NULL) && (::wcscmp (m_bstrKey, bstrNewVal) == 0) )
		return S_OK ;

	// Save the string.
	BSTR	bstrKey = ::SysAllocString (bstrNewVal) ;

	if ( bstrKey == NULL )
	{
		_Error (OLESTR("Key: Out of memory.")) ;
		return E_OUTOFMEMORY ;
	}

	// Free the old key.
	::SysFreeString (m_bstrKey) ;

	// Set the new key.
	m_bstrKey = bstrKey ;

	m_bTouched = TRUE ;
	return S_OK;
}

/*!
 * This method sets an initial vector string.
 *
 * \param[in] bstrNewVal	An initial vector string which is converted into UTF-8 and hashed by MD5.
 *                          Null or an empty string is not allowed.
 */
STDMETHODIMP CEpart::put_InitialVector(const BSTR bstrNewVal)
{
	if ( bstrNewVal == NULL )
	{
		_Error (OLESTR("InitialVector: The initial vector can not be null.")) ;
		return E_INVALIDARG ;
	}
	
	if ( ::SysStringLen (bstrNewVal) == 0 )
	{
		_Error (OLESTR("InitialVector: The initial vector can not be an empty string.")) ;
		return E_INVALIDARG ;
	}

	// If the new initial vector and the old initial vector are the same
	if ( (m_bstrInitialVector != NULL) && (::wcscmp (m_bstrInitialVector, bstrNewVal) == 0) )
		return S_OK ;

	// Save the string.
	BSTR	bstrInitialVector = ::SysAllocString (bstrNewVal) ;

	if ( bstrInitialVector == NULL )
	{
		_Error (OLESTR("InitialVector: Out of memory.")) ;
		return E_OUTOFMEMORY ;
	}

	// Free the old initial vector.
	::SysFreeString (m_bstrInitialVector) ;

	// Set the new initial vector.
	m_bstrInitialVector = bstrInitialVector ;

	m_bTouched = TRUE ;
	return S_OK;
}

/*!
 * This method Initializes internal states.
 */
HRESULT CEpart::_Initialize (void)
	throw ()
{
	if ( !m_bTouched )
		return S_OK ;

	// If the key is NULL
	if ( m_bstrKey == NULL )
	{
		_Error (OLESTR("_Initialize: The key can not be NULL.")) ;
		return E_FAIL ;
	}
	
	// If the initialVector is NULL
	if ( m_bstrInitialVector == NULL )
	{
		_Error (OLESTR("_Initialize: The initialVector can not be NULL.")) ;
		return E_FAIL ;
	}

	HRESULT			hResult ;
	BYTE *			pbyUtf8 = NULL ;
	int				cbyUtf8 ;

	__try
	{
		// Get a UTF-8 byte array from the unicode string.		
		if ( FAILED (hResult = _ToUtf8 (m_bstrKey, &pbyUtf8, &cbyUtf8)) )
			return hResult ;

		// Get the MD5 hash.
		if ( FAILED (hResult = _ToMd5 (pbyUtf8, cbyUtf8, m_pbyHashedKey)) )
			return hResult ;
		
		SAFEFREE (pbyUtf8) ;

		// Get a UTF-8 byte array from the unicode string.
		if ( FAILED (hResult = _ToUtf8 (m_bstrInitialVector, &pbyUtf8, &cbyUtf8)) )
			return hResult ;

		// Get the MD5 hash.
		if ( FAILED (hResult = _ToMd5 (pbyUtf8, cbyUtf8, m_pbyHashedInitialVector)) )
			return hResult ;		
	}
	__finally
	{
		SAFEFREE (pbyUtf8) ;
	}

	return S_OK ;
}

STDMETHODIMP CEpart::encrypt(const BSTR bstrValue, BSTR* pbstrReturn)
{
	if ( pbstrReturn == NULL )
	{
		_Error (OLESTR("Encrypt: The pbstrReturn can not be null.")) ;
		return E_INVALIDARG ;
	}
	*pbstrReturn = NULL ;


	HRESULT		hResult ;

	// Initialize
	if ( FAILED (hResult = _Initialize ()) )
		return hResult ;


	BYTE *		pbyUtf8 = NULL ;
	BYTE *		pbyPkcs7 = NULL ;
	BYTE *		pbyEncrypted = NULL ;
	int			cbyUtf8 = 0 ;
	int			cbyPkcs7 = 0 ;

	
	__try
	{
	
		// Get the UTF-8 byte array.
		if ( FAILED (hResult = _ToUtf8 (bstrValue, &pbyUtf8, &cbyUtf8)) )
			return hResult ;


		// Get the PKCS #7 padded array.
		if ( FAILED (hResult = _ToPkcs7 (pbyUtf8, cbyUtf8, &pbyPkcs7, &cbyPkcs7)) )
			return hResult ;


		// Copy the initial vector.
		BYTE	pbyInitialVector[STRENCRYPTER_BLOCK_SIZE] ;
		::memcpy (pbyInitialVector, m_pbyHashedInitialVector, STRENCRYPTER_BLOCK_SIZE) ;

	
		// Create an output buffer.
		pbyEncrypted = (BYTE *) ::malloc (sizeof (BYTE) * cbyPkcs7) ;
		if ( pbyEncrypted == NULL )
		{
			_Error (OLESTR("Encrypt: Out of memory")) ;
			return E_OUTOFMEMORY ;
		}


		// Set the key to a cryptograph context.
		::aes_context		ctx ;	
		::aes_setkey_enc (&ctx, m_pbyHashedKey, STRENCRYPTER_KEY_BIT_SIZE) ;
	

		// Encrypt data
		::aes_crypt_cbc (&ctx, AES_ENCRYPT, cbyPkcs7, pbyInitialVector, pbyPkcs7, pbyEncrypted) ;
			
		// Get a string encoded by Base64
		if ( FAILED (hResult = _ToBase64 (pbyEncrypted, cbyPkcs7, pbstrReturn)) )
			return hResult ;	
	}
	__finally
	{
		SAFEFREE (pbyUtf8) ;
		SAFEFREE (pbyPkcs7) ;
		SAFEFREE (pbyEncrypted) ;
	}
	
	return S_OK;
}

/*!
 * This method decrypts a string which is encrypted with the same key and initial vector.
 *
 * \param[in] bstrValue		A string to decrypt. It must be a string encrypted with the same key and initial vector.
 *                          Null or an empty string is not allowed.
 * \param[out] pbstrReturn	A decrypted string
 */
STDMETHODIMP CEpart::decrypt(const BSTR bstrValue, BSTR* pbstrReturn)
{
	if ( pbstrReturn == NULL )
	{
		_Error (OLESTR("Decrypt: The pbstrReturn can not be null.")) ;
		return E_INVALIDARG ;
	}
	*pbstrReturn = NULL ;

	if ( bstrValue == NULL )
	{
		_Error (OLESTR("Decrypt: The cipher string can not be null.")) ;
		return E_INVALIDARG ;
	}

	if ( ::SysStringLen (bstrValue) == 0 )
	{
		_Error (OLESTR("Decrypt: The cipher string can not be an empty string.")) ;
		return E_INVALIDARG ;
	}


	HRESULT		hResult ;

	// Initialize
	if ( FAILED (hResult = _Initialize ()) )
		return hResult ;


	BYTE *		pbyEncrypted = NULL ;
	BYTE *		pbyDecrypted = NULL ;
	int			cbyEncrypted = 0 ;
	int			cbyUtf8 = 0 ;

	__try
	{
		// Get the original byte array from a Base64 string.
		if ( FAILED (hResult = _FromBase64 (bstrValue, &pbyEncrypted, &cbyEncrypted)) )
			return hResult ;

		// If the length of the byte array is not a multiple of BLOCK_SIZE
		if ( (cbyEncrypted % STRENCRYPTER_BLOCK_SIZE) > 0 )
		{
			_Error (OLESTR("Decrypt: The length of a cipher data must be a multiple of block size.")) ;
			return E_INVALIDARG ;
		}

		// Copy the initial vector.
		BYTE	pbyInitialVector[STRENCRYPTER_BLOCK_SIZE] ;
		::memcpy (pbyInitialVector, m_pbyHashedInitialVector, STRENCRYPTER_BLOCK_SIZE) ;


		// Create an output buffer.
		pbyDecrypted = (BYTE *) ::malloc (sizeof (BYTE) * cbyEncrypted) ;
		if ( pbyEncrypted == NULL )
		{
			_Error (OLESTR("Decrypt: Out of memory.")) ;
			return E_OUTOFMEMORY ;
		}


		// Set the key to a cryptograph context.
		::aes_context		ctx ;
		::aes_setkey_dec (&ctx, m_pbyHashedKey, STRENCRYPTER_KEY_BIT_SIZE) ;


		// Decrypt data
		::aes_crypt_cbc (&ctx, AES_DECRYPT, cbyEncrypted, pbyInitialVector, pbyEncrypted, pbyDecrypted) ;


		// Gets the actual size of data encoded by PKCS7 padding algorithm.
		if ( FAILED (hResult = _GetActualSizeOfPkcs7 (pbyDecrypted, cbyEncrypted, &cbyUtf8)) )
			return hResult ;


		// Get an unicode string from a UTF-8 string.
		if ( FAILED (hResult = _FromUtf8 (pbyDecrypted, cbyUtf8, pbstrReturn)) )
			return hResult ;
	}
	__finally
	{
		SAFEFREE (pbyDecrypted) ;
		SAFEFREE (pbyEncrypted) ;
	}

	return S_OK;
}

/*!
 * This method reports error information.
 *
 * \param[in] szMessage		An error message.
 */
void CEpart::_Error (LPCOLESTR szMessage)
	throw ()
{
	Error (szMessage, IID_IEpart) ;
}

/*!
 * This method reports error information.
 *
 * \param[in] szMessage		An error message.
 * \param[in] dwErrCode		A WIN32 error code.
 */
void CEpart::_Error (LPCOLESTR szMessage, DWORD dwErrCode)
	throw ()
{
	const int		BUFF_SIZE = 1024 ;

	// Allocate buffer
	OLECHAR		szBuffer[BUFF_SIZE] ;
	szBuffer[0] = NULL ;

	// Get the system error message
	//Ryeol::GetSystemErrMsgW (szBuffer, BUFF_SIZE, dwErrCode, szMessage) ;

	// Set error message
	Error (szBuffer, IID_IEpart) ;
}

/*!
 * This method converts an unicode string into a UTF-8 byte array.
 *
 * \param[in] bstrValue		A string to encode. If this value is NULL, *ppbyOutput is set to NULL
 *                          and *pcbyOutput is set to 0.
 * \param[out] ppbyOutput	A pointer to receive the converted UTF-8 byte array.
 *                          It will point to memory allocated by ::malloc if this method is succeed.
 * \param[out] pcbyOutput	A pointer to receive the length of the UTF-8 byte array.
 *                          If this method is failed, zero is returned.
 */
HRESULT CEpart::_ToUtf8 (const BSTR bstrValue, BYTE ** ppbyOutput, int * pcbyOutput)
	throw ()
{
	ATLASSERT ( ppbyOutput != NULL ) ;
	ATLASSERT ( pcbyOutput != NULL ) ;


	// Initialize variables
	*ppbyOutput = NULL ;
	*pcbyOutput = 0 ;

	if ( bstrValue == NULL )
		return S_OK ;


	int		cchValue = ::SysStringLen (bstrValue) ;

	if ( cchValue == 0 )
		return S_OK ;


	// Get a required buffer size
	*pcbyOutput = ::WideCharToMultiByte (
			CP_UTF8									// UTF-8
			, 0										// Not used if it is UTF-8.
			, reinterpret_cast<LPWSTR> (bstrValue)	// A unicode string.
			, cchValue								// The length of a unicode string. (without NULL)
			, NULL									// Output buffer.
			, 0										// Output buffer size.
			, NULL									// It is not used for UTF-8.
			, NULL									// It is not used for UTF-8.
		) ;

	if ( *pcbyOutput == 0 )
	{
		_Error (OLESTR("_ToUtf8: WideCharToMultiByte - "), ::GetLastError ()) ;
		return E_FAIL ;
	}


	// Allocate memory.
	*ppbyOutput = (BYTE *) ::malloc (*pcbyOutput) ;

	if ( *ppbyOutput == NULL )
	{
		_Error (OLESTR("_ToUtf8: Out of memory.")) ;
		return E_OUTOFMEMORY ;
	}


	// Get the UTF-8 string.
	if ( 0 == ::WideCharToMultiByte (
			CP_UTF8
			, 0
			, reinterpret_cast<LPWSTR> (bstrValue)
			, cchValue	
			, reinterpret_cast<CHAR *> (*ppbyOutput)
			, *pcbyOutput
			, NULL
			, NULL
		) )
	{
		SAFEFREE (*ppbyOutput) ;
		_Error (OLESTR("_ToUtf8: WideCharToMultiByte - "), ::GetLastError ()) ;
		return E_FAIL ;
	}

	return S_OK ;
}

/*!
 * This method converts a UTF-8 byte array into an unicode string.
 *
 * \param[in] pbyUtf8		A UTF-8 byte array.
 * \param[in] cbyUtf8		The length of a UTF-8 byte array.
 *                          If this value is 0, pbstrOutput will point to an empty string allocated by ::SysAllocString.
 * \param[out] pbstrOutput	A pointer to receive the converted unicode string.
 *                          It will point to memory allocated by ::SysAllocString if this method is succeed.
 */
HRESULT CEpart::_FromUtf8 (const BYTE * pbyUtf8, const int cbyUtf8, BSTR * pbstrOutput)
	throw ()
{
	ATLASSERT ( cbyUtf8 >= 0 ) ;

	if ( cbyUtf8 == 0 )
	{
		// If the length of the UTF-8 byte array is zero
		*pbstrOutput = ::SysAllocString (OLESTR("")) ;

		if ( *pbstrOutput == NULL )
		{
			_Error (OLESTR("_FromUtf8: Out of memory.")) ;
			return E_OUTOFMEMORY ;
		}

		return S_OK ;
	}

	ATLASSERT ( pbyUtf8 != NULL ) ;


	// Get the required buffer size.
	int			cchRequired ;
	cchRequired = ::MultiByteToWideChar (
			CP_UTF8														// UTF-8
			, 0															// dwFlags is not used.
			, reinterpret_cast<LPSTR> (const_cast<BYTE *> (pbyUtf8))	// UTF-8 byte array.
			, cbyUtf8													// The length of UTF-8 byte array.
			, NULL														// Buffer.
			, 0															// Buffer size.
	) ;

	if ( 0 == cchRequired )
	{
		_Error (OLESTR("_FromUtf8: MultiByteToWideChar - "), ::GetLastError ()) ;
		return E_FAIL ;
	}

	// Allocate memory.
	*pbstrOutput = ::SysAllocStringLen (NULL, cchRequired) ;

	if ( *pbstrOutput == NULL )
	{
		_Error (OLESTR("_FromUtf8: Out of memory.")) ;
		return E_OUTOFMEMORY ;
	}


	if ( 0 == ::MultiByteToWideChar (
			CP_UTF8														// UTF-8
			, 0															// dwFlags is not used.
			, reinterpret_cast<LPSTR> (const_cast<BYTE *> (pbyUtf8))	// UTF-8 byte array.
			, cbyUtf8													// The length of UTF-8 byte array.
			, *pbstrOutput												// Buffer.
			, cchRequired												// Buffer size.
		) )
	{
		_Error (OLESTR("_FromUtf8: MultiByteToWideChar - "), ::GetLastError ()) ;
		return E_FAIL ;
	}

	return S_OK ;
}

/*!
 * This method returns the MD5 hash.
 *
 * \param[in] pbyValue		A byte array to calculate MD5 hash. NULL is not allowed.
 * \param[in] cbyValue		The length of the byte array. 0 is not allowed.
 * \param[out] pbyBuffer	An allocated buffer to receive MD5 hash. 16bytes buffer is required.
 */
HRESULT CEpart::_ToMd5 (const BYTE * pbyValue, int cbyValue, BYTE * pbyBuffer)
	throw ()
{
	ATLASSERT ( pbyValue != NULL ) ;
	ATLASSERT ( cbyValue >= 0 ) ;
	ATLASSERT ( pbyBuffer != NULL ) ;

	// Get MD5.
	::md5 (const_cast<BYTE *> (pbyValue), cbyValue, pbyBuffer) ;

	return S_OK ;
}

/*!
 * This method encodes data according to the PKCS7 padding algorithm.
 *
 * \param[in] pbyValue		A byte array to pad.
 * \param[in] cbyValue		Length of the byte array. If this value is 0, Only padding is appended.
 * \param[out] ppbyOutput	A pointer to receive the padded byte array.
 *                          It will point memory allocated by ::malloc if this method is succeed.
 * \param[out] pcbyOutput	A pointer to receive the length of the padded byte array.
 */
HRESULT CEpart::_ToPkcs7 (const BYTE * pbyValue, int cbyValue, BYTE ** ppbyOutput, int * pcbyOutput)
	throw ()
{
	ATLASSERT ( cbyValue >= 0 ) ;
	ATLASSERT ( ppbyOutput != NULL ) ;
	ATLASSERT ( pcbyOutput != NULL ) ;

	if ( cbyValue > 0 )
		ATLASSERT ( pbyValue != NULL ) ;


	// Get a number of bytes to pad.
	int		padSize = STRENCRYPTER_BLOCK_SIZE - (cbyValue % STRENCRYPTER_BLOCK_SIZE) ;


	// Allocate memory
	*pcbyOutput = cbyValue + padSize ;
	*ppbyOutput = (BYTE *) ::malloc (sizeof (BYTE) * (*pcbyOutput)) ;

	if ( *ppbyOutput == NULL )
	{
		_Error (OLESTR("_ToPkcs7: Out of memory.")) ;
		return E_OUTOFMEMORY ;
	}


	// Copy the original value
	::memcpy (*ppbyOutput, pbyValue, cbyValue) ;


	// Add padding
	for (int i = 0; i < padSize; i++)
		(*ppbyOutput)[cbyValue + i] = padSize ;


	return S_OK ;
}

/*!
 * This method returns an actual size of data encoded by PKCS7 padding algorithm.
 *
 * \param[in] pbyPkcs7		A byte array in PKCS7 padding algorithm. NULL is not allowed.
 * \param[in] cbyPkcs7		The length of the byte array. cbyPkcs7 must be a multiple of block size.
 * \param[out] piSize		A pointer to receive the actual size.
 */
HRESULT CEpart::_GetActualSizeOfPkcs7 (const BYTE * pbyPkcs7, int cbyPkcs7, int * piSize)
	throw ()
{
	ATLASSERT ( pbyPkcs7 != NULL ) ;
	ATLASSERT ( cbyPkcs7 >= 0 ) ;
	ATLASSERT ( piSize != NULL ) ;


	// The length must be a multiple of block size.
	if ( cbyPkcs7 % STRENCRYPTER_BLOCK_SIZE > 0 )
	{
		_Error (OLESTR("_GetActualSizeOfPkcs7: The length of data is not a multiple of block size.")) ;
		return E_FAIL ;
	}


	// Get the padding size.
	int		padSize = pbyPkcs7[cbyPkcs7 - 1] ;


	// The padding size must be a number greater than 0 and less equal than the block size.
	if ( padSize < 1 || padSize > STRENCRYPTER_BLOCK_SIZE )
	{
		_Error (OLESTR("_GetActualSizeOfPkcs7: The padding size must be a number greater than 0 and less equal than the block size.")) ;
		return E_FAIL ;
	}


	// Check padding.
	for (int i = 0; i < padSize; i++)
	{
		if ( pbyPkcs7[cbyPkcs7 - i - 1] != padSize )
		{
			_Error (OLESTR("_GetActualSizeOfPkcs7: A padded value is not valid.")) ;
			return E_FAIL ;
		}
	}

	
	// Save the size of data without padding.
	*piSize = cbyPkcs7 - padSize ;


	return S_OK ;
}

// Base64 padding character
static const OLECHAR	g_chBase64Padding = '=' ;

// Base64 characters map
static LPCOLESTR		g_pszBase64Characters = 
		OLESTR("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/") ;

/*!
 * This method converts a binary data into a Base64 encoded string.
 *
 * \param[in] pbyValue		A byte array to encode by Base64.
 * \param[in] cbyValue		The length of the byte array.
 *                          If this value is 0, pbstrOutput will point to an empty string allocated by ::SysAllocString.
 * \param[out] pbstrOutput	A pointer to receive the converted unicode string.
 *                          It will point to memory allocated by ::SysAllocString if this method is succeed.
 */
HRESULT CEpart::_ToBase64 (const BYTE * pbyValue, int cbyValue, BSTR * pbstrOutput)
	throw ()
{
	ATLASSERT ( cbyValue >= 0 ) ;
	ATLASSERT ( pbstrOutput != NULL ) ;

	if ( cbyValue > 0 )
		ATLASSERT ( pbyValue != NULL ) ;


	// If the length of the data array is 0, an empty string is returned.
	if ( cbyValue == 0 )
	{
		*pbstrOutput = ::SysAllocString (OLESTR("")) ;
		if ( *pbstrOutput == NULL )
		{
			_Error (OLESTR("_ToBase64: Out of memory.")) ;
			return E_OUTOFMEMORY ;
		}

		return S_OK ;
	}


	// Get a required number of characters.
	int		cchRequired = (cbyValue / 3) * 4 ;
	if ( cbyValue % 3 )
		cchRequired += 4 ;


	// Allocate memory.
	*pbstrOutput = ::SysAllocStringLen (NULL, cchRequired) ;
	if ( *pbstrOutput == NULL )
	{
		_Error (OLESTR("_ToBase64: Out of memory.")) ;
		return E_OUTOFMEMORY ;
	}


	// Converts the byte array into a Base64 string. (3 bytes to 4 characters)
	const int	cbyBuffer = 3 ;

	BYTE		pbyBuffer[cbyBuffer] ;
	int			iBufferIndex = 0, iOutputIndex = 0, iBase64Index ;

	::memset (pbyBuffer, 0, cbyBuffer) ;

	for (int i = 0; i < cbyValue; i++)
	{
		// Accumulate bytes in buffer.
		pbyBuffer[iBufferIndex++] = pbyValue[i] ;

		// If 3 bytes are collected or the last byte is collected
		if ( (iBufferIndex == 3) || (i == cbyValue - 1) )
		{
			// Make a first character
			iBase64Index = (pbyBuffer[0] & 0xFC) >> 2 ;
			(*pbstrOutput)[iOutputIndex++]= g_pszBase64Characters[iBase64Index] ;


			// Make a second character
			iBase64Index = ((pbyBuffer[0] & 0x03) << 4) | (pbyBuffer[1] >> 4) ;
			(*pbstrOutput)[iOutputIndex++] = g_pszBase64Characters[iBase64Index] ;


			// Make a third character
			iBase64Index = ((pbyBuffer[1] & 0x0F) << 2) | (pbyBuffer[2] >> 6) ;
			(*pbstrOutput)[iOutputIndex++] = 
				iBufferIndex > 1 ? g_pszBase64Characters[iBase64Index] : g_chBase64Padding ;


			// Make a fourth character
			iBase64Index = pbyBuffer[2] & 0x3F ;
			(*pbstrOutput)[iOutputIndex++] = 
				iBufferIndex > 2 ? g_pszBase64Characters[iBase64Index] : g_chBase64Padding ;


			// Reset variables related to the buffer
			iBufferIndex = 0 ;
			::memset (pbyBuffer, 0, cbyBuffer) ;
		}
	}

	(*pbstrOutput)[cchRequired] = NULL ;

	return S_OK ;
}


/*!
 * This method converts a Base64 encoded string into a binary data.
 *
 * \param[in] bstrBase64	A Base64 encoded string to decode.
 *                          The length of the string must be a multiple of 4.
 *                          If it is null or an empty string, *ppbyOutput will be NULL and *pcbyOutput will be 0.
 * \param[out] ppbyOutput	A pointer to receive the converted byte array.
 *                          It will point memory allocated by ::malloc if this method is succeed.
 * \param[out] pcbyOutput	A pointer to receive the length of the converted byte array.
 */
HRESULT CEpart::_FromBase64 (const BSTR bstrBase64, BYTE ** ppbyOutput, int * pcbyOutput)
	throw ()
{
	ATLASSERT ( ppbyOutput != NULL ) ;
	ATLASSERT ( pcbyOutput != NULL ) ;

	// If the string is NULL or an empty string
	int		cchBase64 ;
	if ( (bstrBase64 == NULL) || ((cchBase64 = ::SysStringLen (bstrBase64)) == 0) )
	{
		*ppbyOutput = NULL ;
		*pcbyOutput = 0 ;
		return S_OK ;
	}


	// The length of the string must be a multiple of 4.
	if ( cchBase64 % 4 )
	{
		_Error (OLESTR("_FromBase64: The length of a Base64 string must be a multiple of 4.")) ;
		return E_INVALIDARG ;
	}
	

	// Check the Base64 string
	for (int i = 0; i < cchBase64; i++)
	{
		// If the last character is a padding
		if ( (i == cchBase64 - 1) && (bstrBase64[i] == g_chBase64Padding) )
			break ;


		// If the last two characters are padding
		if ( (i == cchBase64 - 2) && (bstrBase64[i] == g_chBase64Padding) )
			if (bstrBase64[i + 1] == g_chBase64Padding)
				break;


		// If it is a Base64 character
		if ( ::wcschr (g_pszBase64Characters, bstrBase64[i]) != NULL )
			continue ;


		_Error (OLESTR("_FromBase64: An invalid character is found in the Base64 string.")) ;
		return E_INVALIDARG ;
	}


	// Get a required number of bytes.
	*pcbyOutput = (cchBase64 / 4) * 3;
	if ( bstrBase64[cchBase64 - 1] == g_chBase64Padding )
		(*pcbyOutput)--;

	if ( bstrBase64[cchBase64 - 2] == g_chBase64Padding )
		(*pcbyOutput)--;


	// Allocate memory.
	*ppbyOutput = (BYTE *) ::malloc (sizeof (BYTE) * *pcbyOutput) ;

	if ( *ppbyOutput == NULL )
	{
		_Error (OLESTR("_FromBase64: Out of memory.")) ;
		return E_OUTOFMEMORY ;
	}


	// Converts the Base64 encoded string into a binary data. (4 characters to 3 bytes)
	const int	cchBuffer = 4 ;

	OLECHAR		pchBuffer[cchBuffer] ;
	int			iBufferIndex = 0, iOutputIndex = 0, iBase64Index ;

	for (int i = 0; i < cchBase64; i++)
	{
		// Accumulate characters in buffer.
		pchBuffer[iBufferIndex++] = bstrBase64[i] ;

		// If 4 characters are collected
		if ( iBufferIndex == cchBuffer )
		{
			// Make a first byte
			iBase64Index = (int) (::wcschr (g_pszBase64Characters, pchBuffer[0]) - g_pszBase64Characters) ;
			(*ppbyOutput)[iOutputIndex] = (BYTE) (iBase64Index << 2) ;

			iBase64Index = (int) (::wcschr (g_pszBase64Characters, pchBuffer[1]) - g_pszBase64Characters) ;
			(*ppbyOutput)[iOutputIndex++] |= iBase64Index >> 4 ;


			// Make a second byte
			if ( pchBuffer[2] == g_chBase64Padding )
				break ;

			(*ppbyOutput)[iOutputIndex] = (BYTE) ((iBase64Index & 0x0F) << 4) ;

			iBase64Index = (int) (::wcschr (g_pszBase64Characters, pchBuffer[2]) - g_pszBase64Characters) ;
			(*ppbyOutput)[iOutputIndex++] |= iBase64Index >> 2 ;


			// Make a third byte
			if ( pchBuffer[3] == g_chBase64Padding )
				break ;

			(*ppbyOutput)[iOutputIndex] = (BYTE) ((iBase64Index & 0x03) << 6) ;

			iBase64Index = (int) (::wcschr (g_pszBase64Characters, pchBuffer[3]) - g_pszBase64Characters) ;
			(*ppbyOutput)[iOutputIndex++] |= iBase64Index ;


			// Reset variables related to the buffer
			iBufferIndex = 0 ;
		}
	}

	return S_OK ;
}

// sha256
string fromDecimal(uint n, int b)
{
	string chars="0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ";
	string result="";
	while(n>0)
	{
		result=chars.at(n%b)+result;
		n/=b;
	}

	return result;
}

uint K[]=
{   0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
	0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
	0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
	0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
	0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
	0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
	0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
	0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2
};

void makeblock(vector<uint>& ret, string p_msg)
{
	uint cur=0;
	int ind=0;
	for(uint i=0; i<p_msg.size(); i++)
	{
		cur = (cur<<8) | (unsigned char)p_msg[i];

		if(i%4==3)
		{
			ret.at(ind++)=cur;
			cur=0;
		}
	}
}

class Block
{
public:
	vector<uint> msg;

	Block():msg(16, 0) { }

	Block(string p_msg):msg(16, 0)
	{ 
		makeblock(msg, p_msg);
	}

};


void split(vector<Block>& blks, string& msg)
{     
	for(uint i=0; i<msg.size(); i+=64)
	{

		try
		{
			makeblock(blks[i/64].msg, msg.substr(i, 64));
		}
		catch(...)
		{
			cout<<"exc";
		}
	}
}


string mynum(uint x)
{
	string ret;
	for(uint i=0; i<4; i++)
		ret+=char(0);

	for(uint i=4; i>=1; i--)	//big endian machine assumed
	{
		ret += ((char*)(&x))[i-1];
	}
	return ret;
}


#define shr(x,n) ((x & 0xFFFFFFFF) >> n)
#define rotr(x,n) (shr(x,n) | (x << (32 - n)))

uint ch(uint x, uint y, uint z)
{
	return (x&y) ^ (~x&z);
}

uint maj(uint x, uint y, uint z)
{
	return (x&y) ^ (y&z) ^ (z&x);
}

uint fn0(uint x)
{
	return rotr(x, 2) ^ rotr(x, 13) ^ rotr(x, 22);
}

uint fn1(uint x)
{
	return rotr(x, 6) ^ rotr(x, 11) ^ rotr(x, 25);
}

uint sigma0(uint x)
{
	return rotr(x, 7) ^ rotr(x, 18) ^ shr(x, 3);
}

uint sigma1(uint x)
{
	return rotr(x, 17) ^ rotr(x, 19) ^ shr(x, 10);
}

STDMETHODIMP CEpart::sha256(const BSTR bstrValue, BSTR* pbstrReturn)
{	
	// wstirng to string
	std::wstring wstr(bstrValue, SysStringLen(bstrValue));
	std::string msg_arr(wstr.length(),' ');
	copy(wstr.begin(), wstr.end(), msg_arr.begin());

	std::string msg;
		
	msg=msg_arr;
	msg_arr += (char)(1<<7);
	uint cur_len = msg.size()*8 + 8;
	uint reqd_len = ((msg.size()*8)/512+1) *512;
	uint pad_len = reqd_len - cur_len - 64;
	
	string pad(pad_len/8, char(0));
	msg_arr += pad;
	string len_str(mynum(msg.size()*8));
	msg_arr = msg_arr + len_str;

	uint num_blk = msg_arr.size()*8/512;
	vector<Block> M(num_blk, Block());
	split(M, msg_arr);

	uint H[]={0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f,  	0x9b05688c, 0x1f83d9ab, 0x5be0cd19};

	for(uint i=0; i<num_blk; i++)
	{             
		vector<uint> W(64, 0);
		for(uint t=0; t<16; t++)
		{
			W[t] = M[i].msg[t];
		}


		for(uint t=16; t<64; t++)
		{
			W[t] = sigma1(W[t-2]) + W[t-7] + sigma0(W[t-15]) + W[t-16];
		}

		uint work[8];
		for(uint i=0; i<8; i++)
			work[i] = H[i];

		for(uint t=0; t<64; t++)
		{

			uint t1, t2;
			t1 = work[7] + fn1(work[4]) + ch(work[4], work[5], work[6]) + K[t] + W[t];
			t2 = fn0(work[0]) + maj(work[0], work[1], work[2]);
			work[7] = work[6];
			work[6] = work[5];
			work[5] = work[4];
			work[4] = work[3] + t1; 
			work[3] = work[2]; 
			work[2] = work[1];
			work[1] = work[0];
			work[0] = t1 + t2;

		}

		for(uint i=0; i<8; i++)
		{
			H[i] = work[i] + H[i];
		}
	}
	string encryptedStrTemp = "";
	for(uint i=0; i<8; i++) {			
		encryptedStrTemp.append(fromDecimal(H[i], 16));
	}
	
	// stirng to wstring
	std::wstring resultWstr(msg_arr.length(),L' ');
	copy(encryptedStrTemp.begin(), encryptedStrTemp.end(), resultWstr.begin());

	// wsting to BSTR
	*pbstrReturn = SysAllocStringLen(resultWstr.data(), resultWstr.size());

	return S_OK;
}
